#include "AST.h"

#include <iostream>
#include <stack>

nmspc::ASTNode::ASTNode(const char& cmd) :cmd(cmd)
{
}

nmspc::OperateASTNode::OperateASTNode(const char& cmd, const std::int64_t& count) :
	nmspc::ASTNode(cmd),
	count(count)
{
}

std::string nmspc::OperateASTNode::generateAssembleCode(const std::int32_t& indentation)
{
	std::string indentation_str;
	for (std::int32_t i = 0; i < indentation; i++)
	{
		indentation_str += "    ";
	}
	std::string assemble_code;
	switch (this->cmd)
	{
	case '+':
		assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
		assemble_code += std::move(indentation_str + "mov bl, byte [rel rax]\n");
		assemble_code += std::move(indentation_str + "add bl, " + std::to_string(this->count) + "\n");
		assemble_code += std::move(indentation_str + "mov byte [rel rax], bl\n");
		break;
	case '-':
		assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
		assemble_code += std::move(indentation_str + "mov bl, byte [rel rax]\n");
		assemble_code += std::move(indentation_str + "sub bl, " + std::to_string(this->count) + "\n");
		assemble_code += std::move(indentation_str + "mov byte [rel rax], bl\n");
		break;
	case '>':
		assemble_code += std::move(indentation_str + "lea rbx, [rel paper_tape]\n");
		assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
		assemble_code += std::move(indentation_str + "sub rax, rbx\n");
		assemble_code += std::move(indentation_str + "add rax, " + std::to_string(this->count) + "\n");
		assemble_code += std::move(indentation_str + "xor rdx, rdx\n");
		assemble_code += std::move(indentation_str + "mov rcx, [rel paper_tape_capacity]\n");
		assemble_code += std::move(indentation_str + "div rcx\n");
		assemble_code += std::move(indentation_str + "mov rax, rdx\n");
		assemble_code += std::move(indentation_str + "add rax, rbx\n");
		assemble_code += std::move(indentation_str + "mov [rel paper_tape_ptr], rax\n");
		break;
	case '<':
		assemble_code += std::move(indentation_str + "lea rbx, [rel paper_tape]\n");
		assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
		assemble_code += std::move(indentation_str + "mov rcx, [rel paper_tape_capacity]\n");
		assemble_code += std::move(indentation_str + "sub rax, rbx\n");
		assemble_code += std::move(indentation_str + "add rax, rcx\n");
		assemble_code += std::move(indentation_str + "sub rax, " + std::to_string(this->count) + "\n");
		assemble_code += std::move(indentation_str + "xor rdx, rdx\n");
		assemble_code += std::move(indentation_str + "div rcx\n");
		assemble_code += std::move(indentation_str + "mov rax, rdx\n");
		assemble_code += std::move(indentation_str + "add rax, rbx\n");
		assemble_code += std::move(indentation_str + "mov [rel paper_tape_ptr], rax\n");
		break;
	case '.':
		for (std::int64_t i = 0; i < this->count; i++)
		{
			assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
			assemble_code += std::move(indentation_str + "lea rcx, [rel fmt]\n");
			assemble_code += std::move(indentation_str + "movzx rdx, byte [rel rax]\n");
			assemble_code += std::move(indentation_str + "call printf\n");
		}
		break;
	case ',':
		for (std::int64_t i = 0; i < this->count; i++)
		{
			assemble_code += std::move(indentation_str + "call _getch\n");
			assemble_code += std::move(indentation_str + "mov bl, al\n");
			assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
			assemble_code += std::move(indentation_str + "mov byte [rel rax], bl\n");
		}
		break;
	default:
		break;
	}
	return assemble_code;
}

std::uint64_t nmspc::LoopASTNode::loop_label_count = -1;

nmspc::LoopASTNode::LoopASTNode() :
	ASTNode(0)
{
	this->loop_label_id = nmspc::LoopASTNode::loop_label_count;
	nmspc::LoopASTNode::loop_label_count++;
}

nmspc::LoopASTNode::~LoopASTNode()
{
	for (nmspc::ASTNode* ast_node : this->ast_nodes)
	{
		delete ast_node;
	}
}

std::string nmspc::LoopASTNode::generateAssembleCode(const std::int32_t& indentation)
{
	std::string indentation_str;
	for (std::int32_t i = 0; i < indentation; i++)
	{
		indentation_str += "    ";
	}
	std::string assemble_code;
	std::string cur_loop_label_start = "loop_" + std::to_string(this->loop_label_id) + "_begin";
	std::string cur_loop_label_end = "loop_" + std::to_string(this->loop_label_id) + "_end";
	std::string pre_indentation_str = indentation_str;
	for (std::int32_t i = 0; i < 4; i++) if (!pre_indentation_str.empty()) pre_indentation_str.pop_back();
	if (indentation > 0)
	{
		assemble_code += std::move(pre_indentation_str + cur_loop_label_start + ":\n");
		assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
		assemble_code += std::move(indentation_str + "mov bl, byte [rel rax]\n");
		assemble_code += std::move(indentation_str + "cmp bl, 0\n");
		assemble_code += std::move(indentation_str + "je " + cur_loop_label_end + "\n");
	}
	for (const auto& node : this->ast_nodes)
	{
		assemble_code += node->generateAssembleCode(indentation + 1);
	}
	if (indentation > 0)
	{
		assemble_code += std::move(indentation_str + "mov rax, [rel paper_tape_ptr]\n");
		assemble_code += std::move(indentation_str + "mov bl, byte [rel rax]\n");
		assemble_code += std::move(indentation_str + "cmp bl, 0\n");
		assemble_code += std::move(indentation_str + "jne " + cur_loop_label_start + "\n");
		assemble_code += std::move(pre_indentation_str + cur_loop_label_end + ":\n");
	}
	return assemble_code;
}

void nmspc::LoopASTNode::addASTNode(nmspc::ASTNode* ast_node)
{
	this->ast_nodes.push_back(ast_node);
}

nmspc::LoopASTNode* nmspc::generateLoopASTNode(const std::string& pbf_code, const std::int64_t& start, const std::int64_t& end)
{
	nmspc::LoopASTNode* loop_ast_node = new nmspc::LoopASTNode();
	std::int64_t i = start;
	while (i <= end)
	{
		char cmd = pbf_code[i];
		if (cmd == '+' || cmd == '-')
		{
			std::int64_t count = 0;
			while (i <= end)
			{
				if (pbf_code[i] == '+')
				{
					count++;
				}
				else if (pbf_code[i] == '-')
				{
					count--;
				}
				else
				{
					break;
				}
				i++;
			}
			if (count == 0)
			{
				continue;
			}
			else if (count > 0)
			{
				loop_ast_node->addASTNode(new nmspc::OperateASTNode('+', count));
			}
			else
			{
				loop_ast_node->addASTNode(new nmspc::OperateASTNode('-', -count));
			}
		}
		else if (cmd == '<' || cmd == '>')
		{
			std::int64_t count = 0;
			while (i <= end)
			{
				if (pbf_code[i] == '>')
				{
					count++;
				}
				else if (pbf_code[i] == '<')
				{
					count--;
				}
				else
				{
					break;
				}
				i++;
			}
			if (count == 0)
			{
				break;
			}
			else if (count > 0)
			{
				loop_ast_node->addASTNode(new nmspc::OperateASTNode('>', count));
			}
			else
			{
				loop_ast_node->addASTNode(new nmspc::OperateASTNode('<', -count));
			}
		}
		else if (cmd == '.')
		{
			std::int64_t count = 0;
			while (i <= end)
			{
				if (pbf_code[i] == '.')
				{
					count++;
				}
				else
				{
					break;
				}
				i++;
			}
			if (count == 0)
			{
				break;
			}
			else if (count > 0)
			{
				loop_ast_node->addASTNode(new nmspc::OperateASTNode(cmd, count));
			}
		}
		else if (cmd == ',')
		{
			std::int64_t count = 0;
			while (i <= end)
			{
				if (pbf_code[i] == ',')
				{
					count++;
				}
				else
				{
					break;
				}
				i++;
			}
			if (count == 0)
			{
				break;
			}
			else if (count > 0)
			{
				loop_ast_node->addASTNode(new nmspc::OperateASTNode(cmd, count));
			}
		}
		else if (cmd == '[')
		{
			std::stack<std::int64_t> st;
			st.push(i);
			std::int64_t j = i + 1;
			for (; j <= end; j++)
			{
				if (pbf_code[j] == ']')
				{
					st.pop();
					if (st.empty())
					{
						break;
					}
				}
				else if (pbf_code[j] == '[')
				{
					st.push(j);
				}
			}
			if (j <= end)
			{
				loop_ast_node->addASTNode(nmspc::generateLoopASTNode(pbf_code, i + 1, j - 1));
				i = j + 1;
			}
			else
			{
				std::cerr << "\033[31m" << "Unexpected cmd! The error occurred at index " + std::to_string(i) << "\033[0m" << std::endl;
				exit(1);
			}
		}
	}
	return loop_ast_node;
}
